#!/usr/bin/env python3

"""
Extract localization message from data record file or dir
and save them into specified  file

Usage:
    extract_path.py  save_fileName  bag|bagfolder

"""

import sys
import os
from cyber.python.cyber_py3 import cyber
from cyber.python.cyber_py3.record import RecordReader
from modules.common_msgs.localization_msgs import localization_pb2

if len(sys.argv) < 3:
    print("Usage: %s <filename> <fbags|bagfolder>" % sys.argv[0])
    sys.exit(0)

filename = sys.argv[1]
fbags = sys.argv[2]

def read_and_write(rf,wf):
    try:
        reader = RecordReader(rf)
    except:
        print("Ignore FILE:", rf)
    for msg in reader.read_messages():
        if msg.topic == "/apollo/localization/pose":
            localization = localization_pb2.LocalizationEstimate()
            localization.ParseFromString(msg.message)
            x = localization.pose.position.x
            y = localization.pose.position.y
            wf.write(str(x) + "," + str(y) + "\n")

with open(filename, 'w') as f:
    if os.path.isfile(fbags):
        read_and_write(fbags,f)
    elif os.path.isdir(fbags):
        for ff in sorted(os.listdir(fbags)):
            read_and_write(os.path.join(fbags,ff),f)
    else:
        raise Exception("Bags file not exist:", fbags)
    
print("XY written to FILE: %s" % filename)
